#coding: --utf-8--

import os
import numpy as np
from random import shuffle

class LoadLabels():
    '''
    说明
    ----
    用于获取图像的标签
    '''
    def __init__(self, label_file, set_root_dir, start=0, end=-1):
        '''
        说明
        ----
        读取标签的初始化函数

        参数
        ----
        - label_file 标签文件
        - set_root_dir 数据集根目录
        - start 获取数据集的开始位置
        - end 获取数据集的结束位置，-1表示直到末尾
        '''
        if start > end and end > 0:
            raise Exception('error parameter! start is supposed to less than end.')

        self.label_file = open(label_file)  # 标签文件
        self.set_root_dir = set_root_dir    # 数据集根目录
        self.label = []
        _start, _end = 0, 0
        is_done, image_name, faces_counts, label_list = self._load()
        while is_done:
            if _start >= start:
                self.label.append([image_name, faces_counts, label_list])
            else:
                _start += 1
            is_done, image_name, faces_counts, label_list = self._load()
            if end < 0 or _end < end:
                _end += 1
            else:
                break
 
        self._MAXSIZE = len(self.label)
        self._index = 0
        print('data set Size: ', self._MAXSIZE)

    def shuffleData(self):
        shuffle(self.label)
        self._index = 0

    def load(self):
        if self._index >= self._MAXSIZE:
            return False
        self.image_name = self.label[self._index][0]
        self.faces_counts = self.label[self._index][1]
        self.label_list = self.label[self._index][2]
        self._index += 1
        if self._index > self._MAXSIZE:
            return False
        return True

    # 开始读取一张图片的label
    def _load(self):
        image_name = self.label_file.readline()
        if not image_name:
            print('done.')
            return False, None, None, None
        self.image_name = (self.set_root_dir + image_name).split('\n')[0]   # 图像文件名
        self.faces_counts = int(self.label_file.readline().split('\n')[0]) # 该图片中的人脸数
        self.label_list = []
        for index in range(self.faces_counts):
            label_str = self.label_file.readline()
            label = [int(x) for x in label_str.split(' ')[:-1]]
            self.label_list.append(label)
        self.label_list = np.array(self.label_list)
        return True, self.image_name, self.faces_counts, self.label_list    
        
    # 获取图片名
    def getFileName(self):
        return self.image_name

    # 获取当前图像中的人脸数
    def getFacesCounts(self):
        return self.faces_counts
    
    # 将当前图像中的人脸按照一个list全部返回
    # [[x1, y1, w, h, blur, expression, illumination, invalid, occlusion, pose],]
    def getLabelsList(self):
        return self.label_list

    # 将当前图像中的人脸按照一个list全部返回
    # [[x1, y1, w, h],]
    def getBoundingBoxesList(self):
        return self.label_list[:, :4]

    # 将当前图像中的人脸按照一个list全部返回
    # [[x1, y1, w, h],]
    # 用于图像缩放时 
    def getBoundingBoxesListX(self, x=1):
        return x*np.array(self.label_list[:, :4])

    # 将当前图像中的人脸的Blur按照一个list全部返回
    def getBlurList(self):
        return self.label_list[:, 4]

    # 将当前图像中的人脸的Expression按照一个list全部返回
    def getExpressionList(self):
        return self.label_list[:, 5]

    # 将当前图像中的人脸的Illumination按照一个list全部返回
    def getIlluminationList(self):
        return self.label_list[:, 6]

    # 将当前图像中的人脸的Occlusion按照一个list全部返回
    def getOcclusionList(self):
        return self.label_list[:, 7]

    # 将当前图像中的人脸的Pose按照一个list全部返回
    def getPoseList(self):
        return self.label_list[:, 8]

    # 将当前图像中的人脸的Invalid按照一个list全部返回
    def getInvalidList(self):
        return self.label_list[:, 9]


if __name__ == "__main__":
    label_file = "/home/jerry/dataSet/faces/wider_face_split/wider_face_train_bbx_gt.txt"
    set_root_dir = "/home/jerry/dataSet/faces/WIDER_train/"
    label = LoadLabels(label_file, set_root_dir)
    # for i in range(10):
    while True:
        if not label.load():
            break
        # print(label.getFileName())
        # print(label.getLabelsList())
        # print(label.getBoundingBoxesList())
        # print(label.getBlurList())
        # print(label.getExpressionList())
        print(label.getFacesCounts())
        # print(label.getInvalidList())